import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


plt.rcParams['figure.figsize'] = (7, 3.5)
plt.rcParams['figure.dpi'] = 150
plt.rcParams['axes.unicode_minus'] = False  #解决坐标轴负数的铅显示问题

"""
leaky_relu(x) = max(\aplha\timesx, x)

- 优点：
    - 避免梯度消失的问题
    - 计算简单
    - 针对Relu函数中存在的Dead Relu Problem，Leaky Relu函数在输入为负值时，给予输入值一个很小的斜率，在解决了负输入情况下的0梯度问题的基础上，也很好的缓解了Dead Relu问题

- 缺点：
    - 输出不是zero-centered
    - ReLU不会对数据做幅度压缩，所以数据的幅度会随着模型层数的增加不断扩张
    - 理论上来说，该函数具有比Relu函数更好的效果，但是大量的实践证明，其效果不稳定，故实际中该函数的应用并不多。
    - 由于在不同区间应用的不同的函数所带来的不一致结果，将导致无法为正负输入值提供一致的关系预测。

"""

#! 1. 定义可以进行前项和后算运算LeakyReLU算子
class MyLeakyReLUAutoGrad(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, alpha):
        ctx.alpha = alpha
        ctx.save_for_backward(x)       # 保存上下文用于反向传播
        return F.leaky_relu(input=x, negative_slope=alpha)
    
    @staticmethod
    def backward(ctx, grad_output):    # y.backward(torch.ones_like(y)) 传进来的
        x,     = ctx.saved_tensors
        ones_x = torch.ones_like(x)
        grad   = torch.where(x < 0, ones_x * ctx.alpha, ones_x)
        return grad, None
    

class LeakyReLU(nn.Module):
    def __init__(self, alpha):
        super(LeakyReLU, self).__init__()
        self.alpha = alpha
    
    def forward(self, x):
        return MyLeakyReLUAutoGrad.apply(x, self.alpha)
    
#! 2. 测试
x   = torch.randn(2, 10, dtype=torch.float32, requires_grad=True)
print(x)

net = LeakyReLU(alpha=0.1)
y   = net(x)
y.backward(torch.ones_like(y))       

print(x.grad)

#! 3. 梯度可视化
# visualization
# inp = torch.arange(-8, 8, 0.05, requires_grad=True)
# out = net(inp)
# out.sum().backward()

# inp_grad = inp.grad

# plt.plot(inp.detach().numpy(),
#          out.detach().numpy(),
#          label=r"$leakyrelu(x)$",
#          alpha=0.7)
# plt.plot(inp.detach().numpy(),
#          inp_grad.numpy(),
#          label=r"$leakyrelu'(x)$",
#          alpha=0.5)
# plt.scatter(0, 0, color='None', marker='o', edgecolors='r', s=50)
# plt.grid()
# plt.legend()
# plt.show()
